import socket
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad


# 密钥和IV
FIXED_KEY = b"MySecretKey12345"
FIXED_IV = b"InitVector123456"


# 解密函数
def decrypt_data(data):
    cipher = AES.new(FIXED_KEY, AES.MODE_CBC, FIXED_IV)
    decrypted_data = unpad(cipher.decrypt(data), AES.block_size)
    return decrypted_data


# 主函数
def start_server():
    # 创建服务器套接字
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    # 服务器设置
    host = 'localhost'
    port = 9999
    server_socket.bind((host, port))
    server_socket.listen(1)
    print(f"服务器已启动，监听地址: {host}:{port}")

    # 主循环
    while True:
        # 等待连接
        client_socket, client_address = server_socket.accept()
        print(f"客户端已连接: {client_address}")

        # 接收文件名
        filename = client_socket.recv(1024).decode('utf-8')
        print(f"接收文件: {filename}")

        # 接收文件大小
        file_size = int(client_socket.recv(1024).decode('utf-8'))

        # 接收加密数据
        encrypted_data = b''
        remaining = file_size

        while remaining > 0:
            chunk = client_socket.recv(min(4096, remaining))
            if not chunk:
                break
            encrypted_data += chunk
            remaining -= len(chunk)

        # 解密并保存
        print("正在解密...")
        decrypted_data = decrypt_data(encrypted_data)

        # 保存文件
        os.makedirs("接收文件", exist_ok=True)
        save_path = os.path.join("接收文件", filename)

        with open(save_path, 'wb') as f:
            f.write(decrypted_data)

        print(f"文件已保存: {save_path}")
        client_socket.send("完成".encode('utf-8'))
        client_socket.close()
        print("连接已关闭")


# 程序入口
if __name__ == "__main__":
    start_server()